import sqlite3

from datetime import date, datetime

from typing import Optional, Union
from langchain_core.tools import tool

# 创建数据库连接
db = "../resource/travel2.sqlite"

@tool
def search_hotels(
    location: Optional[str] = None,
    name: Optional[str] = None,
    price_tier: Optional[str] = None,
    checkin_date: Optional[Union[datetime, date]] = None,
    checkout_date: Optional[Union[datetime, date]] = None,
) -> list[dict]:
    """
    根据位置、名称、价格等级、入住日期和退房日期搜索酒店。

    args：
        location (Optional[str]): 酒店位置。默认为 None。
        name (Optional[str]): 酒店名称。默认为 None。
        price_tier (Optional[str]): 酒店价格等级。默认为 None。示例：中档、中高档、高档、豪华
        checkin_date (Optional[Union[datetime, date]]): 酒店入住日期。默认为 None。
        checkout_date (Optional[Union[datetime, date]]): 酒店退房日期。默认为 None。

    return：
        list[dict]: 符合搜索条件的酒店字典列表。
    """
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    query = "SELECT * FROM hotels WHERE 1=1"
    params = []

    if location:
        query += " AND location LIKE ?"
        params.append(f"%{location}%")
    if name:
        query += " AND name LIKE ?"
        params.append(f"%{name}%")
    if price_tier:
        query += " AND price_tier = ?"
        params.append(price_tier)
    if checkin_date:
        query += " AND checkin_date >= ?"
        params.append(checkin_date)
    if checkout_date:
        query += " AND checkout_date <= ?"
        params.append(checkout_date)

    cursor.execute(query, params)
    results = cursor.fetchall()

    conn.close()

    return [
        dict(zip([column[0] for column in cursor.description], row)) for row in results
    ]


@tool
def book_hotel(hotel_id: int) -> str:
    """
    通过 ID 预订酒店。

    args：
        hotel_id (int)：要预订的酒店 ID。

    return：
        str：一条指示酒店是否预订成功的消息。
    """
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    cursor.execute("UPDATE hotels SET booked = 1 WHERE id = ?", (hotel_id,))
    conn.commit()

    if cursor.rowcount > 0:
        conn.close()
        return f"酒店 {hotel_id} 预订成功。"
    else:
        conn.close()
        return f"未找到 ID 为 {hotel_id} 的酒店。"


@tool
def update_hotel(
    hotel_id: int,
    checkin_date: Optional[Union[datetime, date]] = None,
    checkout_date: Optional[Union[datetime, date]] = None,
) -> str:
    """
    根据酒店的 ID 更新酒店的入住和退房日期。

    args：
        hotel_id (int)：要更新的酒店 ID。
        checkin_date (Optional[Union[datetime, date]])：酒店新的入住日期。默认为 None。
        checkout_date (Optional[Union[datetime, date]])：酒店新的退房日期。默认为 None。

    return：
        str：一条指示酒店是否更新成功的消息。
    """
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    if checkin_date:
        cursor.execute(
            "UPDATE hotels SET checkin_date = ? WHERE id = ?", (checkin_date, hotel_id)
        )
    if checkout_date:
        cursor.execute(
            "UPDATE hotels SET checkout_date = ? WHERE id = ?",
            (checkout_date, hotel_id),
        )

    conn.commit()

    if cursor.rowcount > 0:
        conn.close()
        return f"酒店 {hotel_id} 更新成功。"
    else:
        conn.close()
        return f"未找到 ID 为 {hotel_id} 的酒店。"


@tool
def cancel_hotel(hotel_id: int) -> str:
    """
    通过酒店 ID 取消酒店。

    args：
        hotel_id (int)：要取消的酒店 ID。

    return：
        str：一条消息，指示酒店是否已成功取消。
    """
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    cursor.execute("UPDATE hotels SET booked = 0 WHERE id = ?", (hotel_id,))
    conn.commit()

    if cursor.rowcount > 0:
        conn.close()
        return f"酒店 {hotel_id} 已成功取消。"
    else:
        conn.close()
        return f"未找到 ID 为 {hotel_id} 的酒店。"